import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import os
import time
import pandas as pd
from torch.utils.data import DataLoader, Dataset
from collections import Counter
from sklearn.model_selection import train_test_split

# ---------------------- 超参数配置 ----------------------
class Config:
    """模型配置参数类"""
    device = torch.device('cpu')  # 强制使用CPU，确保无需GPU即可运行
    max_seq_len = 128             # 最大序列长度，根据外卖评论短文本特性设置
    emb_dim = 128                 # 词向量维度，平衡性能与表示能力
    n_head = 4                    # 多头注意力头数，需满足emb_dim % n_head == 0
    dim_feedforward = 256         # 前馈神经网络维度
    num_layers = 2                # Transformer层数，轻量化设计适合CPU
    num_classes = 2               # 分类类别数（外卖评论正负情感二分类）
    batch_size = 64               # 批次大小，根据CPU内存调整
    lr = 0.001                    # 学习率
    epochs = 5                    # 训练轮数，避免过拟合
    data_path = './data/'         # 数据存储路径
    train_file = 'train.txt'      # 训练集文件名
    test_file = 'test.txt'        # 测试集文件名
    csv_file = 'waimai_10k.csv'   # 原始CSV文件名

cfg = Config()

# ---------------------- 工具函数定义 ----------------------
def setup_data_path():
    """创建数据存储目录（如果不存在）"""
    if not os.path.exists(cfg.data_path):
        os.makedirs(cfg.data_path)
        print(f"已创建数据目录: {cfg.data_path}")

def download_and_split_data():
    """下载数据集并分割为训练集和测试集"""
    try:
        # 尝试读取本地CSV文件
        df = pd.read_csv(os.path.join(cfg.data_path, cfg.csv_file))
        print(f"已加载本地数据集: {cfg.csv_file}")
    except FileNotFoundError:
        # 下载数据集
        raise Exception("本地数据文件不存在，开始下载...")

    # 分割数据集
    train_df, test_df = train_test_split(df, test_size=0.1, random_state=42)

    # 保存为txt文件
    train_df.to_csv(os.path.join(cfg.data_path, cfg.train_file), index=False, header=False)
    test_df.to_csv(os.path.join(cfg.data_path, cfg.test_file), index=False, header=False)

    print(f"数据集已分割为: {cfg.train_file} 和 {cfg.test_file}")
    return len(train_df), len(test_df)

# ---------------------- 数据预处理模块 ----------------------
class WaimaiDataset(Dataset):
    """外卖评论数据集"""
    def __init__(self, file_path, tokenizer, vocab, max_len):
        self.data = []
        self.labels = []
        self.tokenizer = tokenizer
        self.vocab = vocab
        self.max_len = max_len

        # 读取数据
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    # 假设格式为"标签,文本"
                    parts = line.strip().split(',', 1)
                    if len(parts) != 2:
                        continue  # 跳过格式错误的行
                    label, text = parts
                    self.labels.append(int(label))
                    self.data.append(text)
                except ValueError:
                    # 处理格式错误的行
                    continue

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        text = self.data[idx]
        label = self.labels[idx]

        # 分词并转换为索引
        tokens = self.tokenizer(text)
        indices = [self.vocab.get(token, self.vocab['<unk>']) for token in tokens]

        # 截断或填充
        if len(indices) < self.max_len:
            indices = indices + [self.vocab['<pad>']] * (self.max_len - len(indices))
        else:
            indices = indices[:self.max_len]

        return torch.tensor(indices), torch.tensor(label)

# 自定义分词器
def chinese_tokenizer(text):
    """中文分词函数（字符级分词）"""
    return list(text)  # 简单的字符级分词


def custom_tokenize(text):
    tokens = []
    current_token = []
    current_type = None

    for char in text:
        # 判断字符类型
        if '\u4e00' <= char <= '\u9fff':
            char_type = 'chinese'
        elif char.isalnum():
            char_type = 'alnum'
        else:
            char_type = 'other'

        # 处理当前字符
        if current_type is None:
            current_type = char_type
            if char_type == 'chinese':
                tokens.append(char)
            else:
                current_token.append(char)
        else:
            if char_type == 'chinese':
                # 中文单独成词
                if current_token:
                    tokens.append(''.join(current_token))
                    current_token = []
                tokens.append(char)
                current_type = None
            elif char_type == 'alnum':
                # 字母数字连续合并
                if current_type == 'alnum':
                    current_token.append(char)
                else:
                    if current_token:
                        tokens.append(''.join(current_token))
                    current_token = [char]
                    current_type = 'alnum'
            else:
                # 其他字符作为分隔符
                if current_token:
                    tokens.append(''.join(current_token))
                    current_token = []
                current_type = None

    # 处理剩余字符
    if current_token:
        tokens.append(''.join(current_token))

    return tokens


def build_vocab(file_path, tokenizer, min_freq=1):
    """构建词汇表（纯Python实现，不依赖torchtext）"""
    counter = Counter()

    # 读取数据并统计词频
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            try:
                _, text = line.strip().split(',', 1)
                tokens = tokenizer(text)
                counter.update(tokens)
            except ValueError:
                continue

    # 过滤低频词并构建词汇表
    vocab = {'<pad>': 0, '<unk>': 1}  # 初始化特殊标记
    idx = 2
    for token, freq in counter.items():
        if freq >= min_freq:
            vocab[token] = idx
            idx += 1

    return vocab

def load_and_preprocess_data():
    """加载并预处理waimai_10k外卖评论数据集"""
    print("开始加载数据...")

    # 确保数据文件存在
    setup_data_path()
    train_path = os.path.join(cfg.data_path, cfg.train_file)
    test_path = os.path.join(cfg.data_path, cfg.test_file)

    if not os.path.exists(train_path) or not os.path.exists(test_path):
        train_size, test_size = download_and_split_data()
    else:
        train_size = sum(1 for _ in open(train_path, 'r', encoding='utf-8'))
        test_size = sum(1 for _ in open(test_path, 'r', encoding='utf-8'))

    print(f"训练集大小: {train_size}, 测试集大小: {test_size}")

    # 使用自定义分词器
    tokenizer = chinese_tokenizer
    # tokenizer = custom_tokenize

    # 构建词汇表
    print("构建词汇表...")
    vocab = build_vocab(train_path, tokenizer)
    print(f"词表大小: {len(vocab)}")

    # 创建数据集
    train_dataset = WaimaiDataset(train_path, tokenizer, vocab, cfg.max_seq_len)
    test_dataset = WaimaiDataset(test_path, tokenizer, vocab, cfg.max_seq_len)

    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=cfg.batch_size,
        shuffle=True,
        num_workers=0  # Windows环境下设置为0避免多进程问题
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=0
    )

    print("数据预处理完成\n")
    return train_loader, test_loader, vocab


class PositionalEncodingV1(nn.Module):
    """位置编码层（为序列添加位置信息）"""

    def __init__(self, d_model, max_len=5000):
        """
        初始化位置编码层。

        Args:
            d_model (int): 模型的特征维度，通常是词向量的维度。
            max_len (int, optional): 最大序列长度，默认为 5000。
        """
        super().__init__()
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.zeros(d_model) + 1
        pe = position * div_term / max_len
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        """
        Args:
            x: 输入张量 (batch_size, seq_len, emb_dim)

        Returns:
            带位置编码的输入张量
        """
        return x + self.pe[:, :x.size(1), :].to(cfg.device) # type: ignore


class PositionalEncoding(nn.Module):
    """位置编码层（为序列添加位置信息）"""

    def __init__(self, d_model, max_len=5000):
        """
        初始化位置编码层。

        Args:
            d_model (int): 模型的特征维度，通常是词向量的维度。
            max_len (int, optional): 最大序列长度，默认为 5000。
        """
        super().__init__()
        # 初始化位置编码矩阵，用于存储每个位置的编码信息
        # 创建一个形状为 (max_len, d_model) 的全零矩阵，后续会填充位置编码值
        pe = torch.zeros(max_len, d_model)
        # 生成从 0 到 max_len-1 的位置序列，并增加一个维度，形状变为 (max_len, 1)
        # 这个序列代表了序列中每个元素的位置
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 计算位置编码的除数项，用于生成正弦和余弦函数的频率
        # 首先生成从 0 开始，步长为 2，到 d_model-1 的序列
        # 然后乘以 -log(10000) / d_model，并取指数，得到不同维度对应的频率
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # 将位置序列与频率项相乘，得到每个位置在不同维度上的角度值
        # 形状为 (max_len, d_model/2)
        position_div_term = position * div_term
        # 偶数位置使用正弦函数进行编码
        # 将正弦函数计算结果填充到位置编码矩阵的偶数索引列
        pe[:, 0::2] = torch.sin(position_div_term)
        # 奇数位置使用余弦函数进行编码
        # 将余弦函数计算结果填充到位置编码矩阵的奇数索引列
        pe[:, 1::2] = torch.cos(position_div_term)
        # 注册为缓冲区（不参与梯度更新），并增加 batch 维度，形状变为 (1, max_len, d_model)
        # 缓冲区可以被模型保存和加载，但不会被当作模型参数进行优化
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        """
        Args:
            x: 输入张量 (batch_size, seq_len, emb_dim)

        Returns:
            带位置编码的输入张量
        """
        return x + self.pe[:, :x.size(1), :].to(cfg.device) # type: ignore


class TransformerClassifier(nn.Module):
    """轻量级Transformer文本分类模型"""

    def __init__(self, vocab_size, emb_dim, num_classes, num_layers, n_head, dim_feedforward):
        """
        初始化轻量级Transformer文本分类模型。

        Args:
            vocab_size (int): 词汇表的大小，即不同词元的数量。
            emb_dim (int): 词向量的维度，用于表示每个词元的特征。
            num_classes (int): 分类的类别数量。
            num_layers (int): Transformer编码器的层数。
            n_head (int): 多头注意力机制的头数。
            dim_feedforward (int): 前馈神经网络的维度。
        """
        super().__init__()
        # 检查多头注意力头数是否整除词向量维度，因为多头注意力机制需要将词向量维度均分到各个头中
        assert emb_dim % n_head == 0, "词向量维度必须能被注意力头数整除"

        # 词嵌入层，将输入的词索引转换为对应的词向量
        self.embedding = nn.Embedding(vocab_size, emb_dim)
        # 位置编码层，为输入序列添加位置信息，帮助模型感知词元在序列中的位置
        # self.pos_encoder = PositionalEncodingV1(emb_dim)
        self.pos_encoder = PositionalEncoding(emb_dim)

        # 定义单个Transformer编码器层
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=emb_dim,  # 模型的输入特征维度，即词向量的维度
            nhead=n_head,  # 多头注意力机制的头数
            dim_feedforward=dim_feedforward,  # 前馈神经网络的维度
            dropout=0.1,  # 丢弃率，用于防止过拟合
            batch_first=True  # 输入格式设置为 (batch, seq, feature)
        )
        # 堆叠多个Transformer编码器层，构建完整的Transformer编码器
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_layers  # Transformer编码器层的数量
        )

        # 分类层（使用全局池化）
        # 全局平均池化，将序列维度压缩为1，得到每个样本的全局特征表示
        self.pooling = nn.AdaptiveAvgPool1d(1)
        # 分类全连接层，将全局特征映射到分类类别空间
        self.fc = nn.Linear(emb_dim, num_classes)

    def forward(self, src):
        """
        Args:
            src: 输入序列 (batch_size, seq_len)

        Returns:
            logits: 分类对数概率 (batch_size, num_classes)
        """
        # 词嵌入 + 缩放
        embedded = self.embedding(src) * math.sqrt(cfg.emb_dim)
        # 添加位置编码
        pos_embedded = self.pos_encoder(embedded)
        # Transformer编码
        encoded = self.transformer_encoder(pos_embedded)  # (batch, seq, emb_dim)
        # 全局平均池化
        pooled = self.pooling(encoded.transpose(1, 2)).squeeze(2)  # (batch, emb_dim)
        # 分类预测
        logits = self.fc(pooled)
        return logits

# ---------------------- 训练与评估模块 ----------------------
def train(model, train_loader, optimizer, criterion):
    """训练模型一个 epoch"""
    model.train()
    epoch_loss = 0
    start_time = time.time()

    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(cfg.device), labels.to(cfg.device)

        # 梯度清零
        optimizer.zero_grad()
        # 前向传播
        logits = model(inputs)
        # 计算损失
        loss = criterion(logits, labels)
        # 反向传播
        loss.backward()
        # 梯度更新
        optimizer.step()

        epoch_loss += loss.item()

        # 打印训练进度
        if (i + 1) % 50 == 0:
            elapsed = time.time() - start_time
            print(f"训练进度: {i+1}/{len(train_loader)}, 批次损失: {loss.item():.4f}, 耗时: {elapsed:.2f}s")

    return epoch_loss / len(train_loader)

def evaluate(model, test_loader, criterion):
    """评估模型在测试集上的性能"""
    model.eval()
    epoch_loss = 0
    correct = 0

    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(cfg.device), labels.to(cfg.device)

            # 前向传播
            logits = model(inputs)
            # 计算损失
            loss = criterion(logits, labels)
            epoch_loss += loss.item()
            # 计算准确率
            pred = logits.argmax(dim=1)
            correct += (pred == labels).sum().item()

    accuracy = correct / len(test_loader.dataset)
    return epoch_loss / len(test_loader), accuracy

# ---------------------- 主程序模块 ----------------------
def main():
    """主函数：模型训练与评估流程"""
    print("=" * 50)
    print("轻量级Transformer文本分类模型 (CPU优化版)")
    print(f"训练配置: 批次大小={cfg.batch_size}, 学习率={cfg.lr}, 训练轮数={cfg.epochs}")
    print("=" * 50 + "\n")

    # 1. 数据准备
    setup_data_path()
    train_loader, test_loader, vocab = load_and_preprocess_data()

    # 2. 模型初始化
    print("初始化模型...")
    model = TransformerClassifier(
        vocab_size=len(vocab),
        emb_dim=cfg.emb_dim,
        num_classes=cfg.num_classes,
        num_layers=cfg.num_layers,
        n_head=cfg.n_head,
        dim_feedforward=cfg.dim_feedforward
    ).to(cfg.device)

    print(f"模型参数总量: {sum(p.numel() for p in model.parameters())}")
    print(model)
    print()

    # 3. 训练配置
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)

    # 4. 训练循环
    print("开始训练...")
    for epoch in range(cfg.epochs):
        start_time = time.time()

        # 训练
        train_loss = train(model, train_loader, optimizer, criterion)

        # 评估
        test_loss, test_acc = evaluate(model, test_loader, criterion)

        # 打印训练结果
        elapsed = time.time() - start_time
        print(f"\nEpoch {epoch+1}/{cfg.epochs} - 耗时: {elapsed:.2f}s")
        print(f"训练损失: {train_loss:.4f} | 测试损失: {test_loss:.4f}")
        print(f"测试准确率: {test_acc:.4f} | 正确样本数: {int(test_acc * len(test_loader.dataset))}/{len(test_loader.dataset)}\n")

    print("训练完成！")
    print("=" * 50)


if __name__ == "__main__":
    main()
